"""
写入 CSV
"""

import pandas as pd

from my_typing import *
from utils.data_utils import ExcelHeaderInfo


class CSVWriter:
    """ CSV 写入器 """

    def __init__(
        self,
        file_name: StrOrBytesIO,
        header_info_list: list[ExcelHeaderInfo],
        data_list: list[dict[str, AnyOrNone]],
        show_header: bool = True,
        encoding: Literal["UTF-8", "GBK"] = "UTF-8",
    ) -> None:
        """
        Args:
            file_name: 文件名/文件流
            header_info_list: 标题信息列表
            data_list: 数据映射列表
            show_header: 是否写入标题行
            encoding: 字符编码（默认为 UTF-8）
        """

        self.header_en_2_cn: DictStrStr = {
            header_info.name: header_info.label
            for header_info in header_info_list
        }
        """ 标题名称映射（英文标题-中文标题） """

        self.header_cn_2_type: dict[str, Callable] = {
            header_info.label: header_info.type
            for header_info in header_info_list
        }
        """ 标题类型映射（中文标题-数据类型） """

        if len(self.header_en_2_cn) != len(data_list[0]):
            raise ValueError("标题长度与内容长度不符")

        if set(self.header_en_2_cn) != set(data_list[0]):
            raise ValueError("标题 Key 与内容 Key 不一致")

        self.data_list: list[DictStrAny] = data_list
        """ 数据映射列表 """

        self.file_name: StrOrBytesIO = file_name
        """ CSV 文件 """

        self.show_header: bool = show_header
        """ 是否写入标题行 """

        self.encoding: str = encoding
        """ 字符编码 """

    def write(self) -> None:
        """ 写入数据 """

        cn_label_2_values: dict[str, list] = {}
        """ 中文标题-值列表映射 """

        for data in self.data_list:
            for en_label, value in data.items():
                cn_label: str = self.header_en_2_cn.get(en_label)
                cn_label_2_values.setdefault(cn_label, []).append(value)

        df: pd.DataFrame = pd.DataFrame(cn_label_2_values)
        df.to_csv(
            path_or_buf=self.file_name,
            index=False,
            header=self.show_header,
            encoding=self.encoding,
        )
